import os
import cv2
import time
import sys
import json
import datetime
import numpy as np
import multiprocessing as mp
from ultralytics import YOLO
from oss.example_oss import OSS
from python_linkkit_examples.mqtt_sub_pub_on import mqtt

# 该文件要放在根目录
ROOT = '/home/court/project'

def log(ip,s):

    print(s)
    now = str(datetime.datetime.now())
    s = now+': '+s+'\n'
    with open(ROOT+'/logs/log_{}.txt'.format(ip),'a') as f:
        f.write(s)
    f.close()

def rpcc_log(s):
    
    print(s)
    now = str(datetime.datetime.now())
    s = now+': '+s+'\n'
    with open(ROOT+'/logs/rrpc_log.txt','a') as f:
        f.write(s)
    f.close()


def GetRpccMsg(pull_queues,push_queues):

    mt = mqtt()
    time.sleep(5) # 需要sleep，有时候会连接出错，这里需要重联机制
    while 1:
        rc, mid = mt.lk.subscribe_topic(mt.lk.to_full_topic("user/get"))
        if rc == 0:
            rpcc_log("subscribe topic success:%r, mid:%r" % (rc, mid))
        else:
            rpcc_log("subscribe topic fail:%d" % rc)
        while 1:
            if mt.rpcc_msg != '': # 接收到消息
                msg_json = json.loads(mt.rpcc_msg) # 转成json
                start_msg = msg_json['instruct']
                if start_msg == 'end_record': # 如果是结束的标识，则重复放一些字符串，给子线程读取
                    for _ in range(16): # 保证充足
                        pull_queues.put(mt.rpcc_msg)
                    rpcc_log('GetRpccMsg: receive rpcc messge: {},queue size: {}'.format(mt.rpcc_msg,pull_queues.qsize()))
                else:
                    pull_queues.put(mt.rpcc_msg)
                    rpcc_log('GetRpccMsg: receive rpcc messge: {},queue size: {}'.format(mt.rpcc_msg,pull_queues.qsize()))
                mt.rpcc_msg = ''
            elif push_queues.qsize() > 0: # 返回启动成功的标识
                pushmsg = push_queues.get()
                rpcc_log('GetRpccMsg: push rpcc messge: {},queue size: {}'.format(pushmsg,push_queues.qsize()))
                mt.response_rrpc(pushmsg)
            else:
                time.sleep(1)


def get_model(camera):

    model_path = {'L2':'/home/court/project/yolov8n.pt','L1':'/home/court/project/yolov8n.pt','TM1':None,'BM1':None,'TM1':None,'TM1':None,'BM1':None,'BM2':None}
    
    if camera == 'L2':
        model = YOLO(model_path[camera])  # 绝对路径
    elif camera == 'L1': #TODO
        model = YOLO(model_path[camera])

    return model

def Inference(pull_queues,push_queues, ip, ref, videoPath):

    myoss = OSS()
    idx_img = 0
    ret = True
    model = get_model(ref)
    cap1 = cv2.VideoCapture("rtsp://%s:%s@%s/h265/ch1/main/av_stream" % ("admin", "playtrac@1", "{}:554/".format(ip)))
    if cap1.isOpened():
        log(ip,' camera is open')

    start_time = time.time()
    # print('Inference: receive rpcc messge: {},queue size: {}'.format(msg,queues.qsize()))
    while ret:

        if pull_queues.qsize() > 0:
            msg = pull_queues.get()
            msg_json = json.loads(msg)
            start_msg = msg_json['instruct']
            cameraIps = msg_json['cameraIps'].split(',')
            if ip in cameraIps and start_msg == 'end_record': # 接受信息，关闭线程
                log(ip,'camera is closed')
                break # do anything

        ret,img = cap1.read()
        if not ret:
            log(ip,'read img error ')
            continue
            
        idx_img = idx_img + 1
        h,w,_ = img.shape
        # img = cv2.resize(img, (int(w / 4), int(h / 4)))
        # model.predict(img, save=True, imgsz=640, conf=0.3, show=False)
        results = model(img)  # list of Results objects
        
        # cv2.imshow(ref, img)
        # cv2.waitKey(1)

    end_time = time.time()
    s = {
            "ip":str(ip),
            "cameraRef":ref,
            "result":"success",
            "videoPath":videoPath,
            "beginTime":str(start_time),
            "duration":str(end_time-start_time),
        }
    push_queues.put(str(s))

    time.sleep(1)
    log(ip,'start process video and upload video,waiting~')
    start_time = time.time()
    srcpath = os.path.join(ROOT,videoPath,ref+'.jpg')
    dstpath = os.path.join(videoPath,ref+'.jpg')
    print('srcpath ',srcpath)
    print('dstpath ',dstpath)
    myoss.uploadfile(dstpath,srcpath)
    end_time = time.time()
    log(ip, 'process video and upload video finished, use time: {}m'.format(str((end_time-start_time)/60)))
    time.sleep(1)


def StartSubProcess(pull_queues,push_queues):

    
    while 1:
        start_msg = ''
        while 1:  # 等待rpcc响应
            if pull_queues.qsize() > 0:
                msg = pull_queues.get()
                msg_json = json.loads(msg)
                start_msg = msg_json['instruct']
                if start_msg == 'begin_record':
                    break
            time.sleep(1)

        # 返回启动成功
        push_queues.put(str({"reuslt":"success"}))
        cameraIps = msg_json['cameraIps'].split(',')
        cameraRef = msg_json['cameraRef'].split(',')
        videoPath = msg_json['videoPath']

        processes = []
        for ip, ref in zip(cameraIps, cameraRef):
            processes.append(mp.Process(target=Inference, args=(pull_queues,push_queues, ip, ref, videoPath)))

       
        for process in processes:
            # process.daemon = True
            process.start()
        for process in processes:
            process.join()

        print(cameraIps,' is closed!')


def run():

    mp.set_start_method(method='spawn')  # init
    pull_queues = mp.Queue(maxsize=16)
    push_queues = mp.Queue(maxsize=16)
    processes = [mp.Process(target=GetRpccMsg, args=(pull_queues,push_queues)),
                 mp.Process(target=StartSubProcess, args=(pull_queues,push_queues)),]

    for process in processes:
        # process.daemon = True
        process.start()
    for process in processes:
        process.join()
    # p.join()



if __name__ == '__main__':
    run()
    
    
